package com.anji.hyperneat.nd;

import java.util.List;

import javax.xml.bind.JAXBException;
import com.anji.topology.Layer;
import com.anji.nn.activationfunction.ActivationFunction;
import com.anji.nn.activationfunction.ActivationFunctionFactory;
import com.anji.topology.XmlUtil;

public class GridNetND implements ActivatorND, Cloneable {
	
    /**
     * base XML tag
     */
    public final static String XML_TAG = "grid network";

    protected ActivationFunction activationFunction;

    protected boolean isFeedForward;

    protected boolean enableBias;
	
    // Activation values for all Layers
    protected NDFloatArray[] layers;
	
	
	/* Connection weights, from one layer to the next 
	 * weights[0] = weights from layer 0 to 1, weights[1] = weights from layer 1 to 2, etc,
	 * since each layer can have it's own shape. 
	 */
    protected NDFloatArray[] weights;
	
    protected NDFloatArray[] bias;

    protected int maxDimensions;
	
    protected String name;

    protected int cyclesPerStep; //Number of activation cycles to perform per step()
    
    /**
     * Create a GridNetND.
     * @param allLayers List containing layer dimensions, as retrieved from a topology xml descriptor.
     * @param weights The weights arrays for this net.
     * @param bias  The bias arrays for this net.
     * @param activationFunction The activation function.
     * @param maxDimensions The maximum number of dimensions for any layer; 
     * 		all layers are "cast" to this size, 
     * 		i.e. given additional size 1 dimensions to match the number of dimensions of the highest dimensional layer.
     * @param cyclesPerStep The number of cycles per step for recurrent nets; currently unused.
     * @param name The name of the net.
     */
    public GridNetND(List<Layer> allLayers
    		, NDFloatArray[] weights
    		, NDFloatArray[] bias
    		, ActivationFunction activationFunction
    		, int maxDimensions
    		, int cyclesPerStep
    		, boolean enableBias
    		, String name) {
    	this.setLayers(new NDFloatArray[allLayers.size()]);
    	for (int l = 0; l < allLayers.size(); l++) {
    		layers[l] = new NDFloatArray(maxDimensions, XmlUtil.convertByteListToArray(allLayers.get(l).getDimensions()));
    	}    	
    	
    	// Sanity checks
    	for (int l = 1; l < getLayers().length; l++) {
//    		int castSize = Math.max(layers[l].getNumDimensions(), layers[l-1].getNumDimensions());
//    		NDFloatArray temp = new NDFloatArray(castSize, layers[l].getDimensions(), layers[l-1].getDimensions());
    		NDFloatArray temp = new NDFloatArray(getLayers()[l].getDimensions(), getLayers()[l-1].getDimensions());
    		if (!weights[l-1].dimensionsMatch(temp))
    			throw new IllegalArgumentException(
    					"Weight matrix " + (l-1) + " does not match the dimensions for layers " + (l-1) + " and " + l 
    					+ " (Provided " 
    					+ weights[l-1].getDimensionString() 
    					+ ", should have been " 
    					+ temp.getDimensionString() 
    					+ ")."
				);
    	}
  
    	for (int l = 1; l < getLayers().length; l++) {
    		if (!bias[l-1].dimensionsMatch(getLayers()[l]))
    			throw new IllegalArgumentException(
    					"Bias matrix " + (l-1) + " does not match the dimensions for layer " + l
    					+ " (Provided " 
    					+ bias[l-1].getDimensionString() 
    					+ ", should have been " 
    					+ getLayers()[l].getDimensionString() 
    					+ ")."
				);
    	}
    	
    	this.weights = weights;
    	this.bias = bias;
    	this.setActivationFunction(activationFunction);
    	this.name = name;
    	this.setMaxDimensions(maxDimensions);
    	this.setCyclesPerStep(cyclesPerStep);
    	this.enableBias = enableBias;
	}

    /**
     * Create a GridNetND.
     * @param layers The actual activation value arrays that the network will use.
     * @param weights The weights arrays for this net.
     * @param bias  The bias arrays for this net.
     * @param activationFunction The activation function.
     * @param maxDimensions The maximum number of dimensions for any layer; 
     * 		all layers are "cast" to this size, 
     * 		i.e. given additional size 1 dimensions to match the number of dimensions of the highest dimensional layer.
     * @param cyclesPerStep The number of cycles per step for recurrent nets; currently unused.
     * @param name The name of the net.
     */
    public GridNetND(NDFloatArray[] layers
    		, NDFloatArray[] weights
    		, NDFloatArray[] bias
    		, ActivationFunction activationFunction
    		, int maxDimensions
    		, int cyclesPerStep
    		, boolean enableBias
    		, String name) {
    	this.setLayers(layers); 	
    	
    	// Sanity checks
    	for (int l = 1; l < layers.length; l++) {
    		NDFloatArray temp = new NDFloatArray(layers[l].getDimensions(), layers[l-1].getDimensions());
    		if (!weights[l-1].dimensionsMatch(temp))
    			throw new IllegalArgumentException(
    					"Weight matrix " + (l-1) + " does not match the dimensions for layers " + (l-1) + " and " + l 
    					+ " (Provided " 
    					+ weights[l-1].getDimensionString() 
    					+ ", should have been " 
    					+ temp.getDimensionString() 
    					+ ")."
				);
    	}
  
    	for (int l = 1; l < layers.length; l++) {
    		if (!bias[l-1].dimensionsMatch(layers[l]))
    			throw new IllegalArgumentException(
    					"Bias matrix " + (l-1) + " does not match the dimensions for layer " + l
    					+ " (Provided " 
    					+ bias[l-1].getDimensionString() 
    					+ ", should have been " 
    					+ layers[l].getDimensionString() 
    					+ ")."
				);
    	}
    	
    	this.weights = weights;
    	this.bias = bias;
    	this.setActivationFunction(activationFunction);
    	this.name = name;
    	this.setMaxDimensions(maxDimensions);
    	this.setCyclesPerStep(cyclesPerStep);
    	this.enableBias = enableBias;
	}
    
	/* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#cost()
	 */
    @Override
	public long cost() {
        return getConnectionCount(true) * getActivationFunction().cost();
    }
    
    
    /* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#getConnectionCount(boolean)
	 */
    @Override
	public int getConnectionCount(boolean includeBias) {
    	int connectionCount = 0;
    	for (int i = 0; i < weights.length; i++) {
    		connectionCount += weights[i].size();
    	}
    	if (includeBias) {
	    	for (int i = 0; i < bias.length; i++) {
	    		connectionCount += bias[i].size();
	    	}
    	}
    	
    	return connectionCount;
    }
    
    
    /* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#setInputs(com.anji.hyperneat.nd.NDFloatArray)
	 */
    @Override
	public void setInputs(NDFloatArray inputs) {
    	if (this.getLayers()[0].dimensionsMatch(inputs))  
    		this.getLayers()[0] = inputs;
    	else {
    		throw new IllegalArgumentException(
    				"Dimensions of inputs do not match network input size (Provided " 
    				+ inputs.getDimensionString() 
    				+ ", should have been " 
    				+ this.getLayers()[0].getDimensionString() + ")." 
			);
    	}
    }
    
    /* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#getInputs()
	 */
    @Override
	public NDFloatArray getInputs() {
        return getLayers()[0];
    }

    /* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#getOutputs()
	 */
    @Override
	public NDFloatArray getOutputs() {
        return getLayers()[getLayers().length-1];
    }

    @Override
	public NDFloatArray getLayer(int layer) {
    	return getLayers()[layer];
    }
    
    @Override
	public int getNumLayers() {
    	return getLayers().length;
    }

	/* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#isFeedForward()
	 */
    @Override
	public boolean isFeedForward() {
        return isFeedForward;
    }
    
    public boolean getEnableBias() {return enableBias;}
    
    /** 3DRF revisit this
     * Provides a reference to the internal weights array. Format is 
     * [z1][y1][x1][zR][yR][xR], where xR, yR, and zR are relative to x1, y1
     * and z1, and with the corresponding negative direction ranges from
     * connectionMaxRanges added to avoid negative array indices.
     * x1, y1, z1 is the target and xR, yR, zR is the source coordinates. Note 
     * that the range of connections for neurons at or close to the edge of 
     * the grid are circumscribed by the edge of the grid, so the last three 
     * dimensions of the array will have differing sizes to account for this 
     * (except in the unlikely scenario that connections extend only straight 
     * down). Note that modifications to the returned array are reflected in 
     * the operation of the GridNet. 
     * @return A reference to the weights array. 
     */
    public NDFloatArray[] getWeights() {
    	return weights;
    }
    
    public NDFloatArray[] getBias() {
    	return bias;
    }
	

	
	public Object next() {
		// TODO Auto-generated method stub
		return null;
	}

	
	/* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#next(com.anji.hyperneat.nd.NDFloatArray)
	 */
	@Override
	public NDFloatArray next(NDFloatArray stimuli) {
		setInputs(stimuli);
		stepFF();
		return getOutputs();																																
	}

	
	/* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#nextSequence(com.anji.hyperneat.nd.NDFloatArray[])
	 */
	@Override
	public NDFloatArray[] nextSequence(NDFloatArray[] stimuli) {
		NDFloatArray[] outs = new NDFloatArray[stimuli.length];
		for (int i = 0; i < outs.length; i++) {
			stepFF();
			outs[i] = getOutputs();
		}			
		
		return outs;
	}


	/* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#reset()
	 */
	@Override
	public void reset() {
		for (NDFloatArray layer : getLayers()) {
			layer.clear();
		}
	}

	

	/* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#getName()
	 */
	@Override
	public String getName() {return name;}
	public float getMinResponse() {return 0;}
	public float getMaxResponse() {return 1;}
	/* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#setName(java.lang.String)
	 */
	@Override
	public void setName(String string) {this.name = string;}
	
	/* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#getInputDimension()
	 */
	@Override
	public int[] getInputDimension() {
		return layers[0].getDimensions();
	}

	
	/* (non-Javadoc)
	 * @see com.anji.hyperneat.nd.ActivatorND#getOutputDimension()
	 */
	@Override
	public int[] getOutputDimension() {
		return layers[layers.length-1].getDimensions();
	}

	

    /**
     * Perform one complete cycle for a feed forward network, propogating signal
     * from input layer to output layer, only activating each layer once in
     * sequence (connections can only exist from layer n to layer n+1).
     */
    public void stepFF() {
        //for each target layer
    	for (int l = 1; l < getLayers().length; l++) {
    		for (NDFloatArray.MatrixIterator tgtIt = getLayers()[l].iterator(); tgtIt.hasNext(); tgtIt.next()) {
    			float sum = enableBias ? bias[l-1].get(tgtIt.getCurrentCoordinates()) : 0;
    			
    			for (NDFloatArray.MatrixIterator srcIt = getLayers()[l-1].iterator(); srcIt.hasNext(); srcIt.next()) {
    				sum += srcIt.get() * weights[l-1].get(tgtIt.getCurrentCoordinates(), srcIt.getCurrentCoordinates());
    			}
    			
    			tgtIt.set( getActivationFunction().apply(sum) );
    		}
    	}
    }



	/**
	 * For Testing only
	 * @param args
	 * @throws JAXBException 
	 */
	public static void main(String[] args) throws JAXBException {
		com.anji.topology.Topology td = XmlUtil.getTopologyDescriptor("test.xml");
		List<Layer> allLayers = XmlUtil.getAllLayers(td);
		NDFloatArray wgts1 = new NDFloatArray(2, 2, 1, 2, 2, 2);
		for (
				com.anji.hyperneat.nd.NDFloatArray.MatrixIterator it = wgts1.iterator();
				it.hasNext();
				it.next()
			) {
			it.set(0.25f);
		}
		NDFloatArray wgts2 = new NDFloatArray(2, 2, 1, 2, 2, 1);
		for (
				com.anji.hyperneat.nd.NDFloatArray.MatrixIterator it = wgts2.iterator();
				it.hasNext();
				it.next()
			) {
			it.set(0.25f);
		}
		
		ActivationFunction activationFn = ActivationFunctionFactory.getInstance().get("sigmoid");
		NDFloatArray bias1 = new NDFloatArray(2, 2, 1);
		for (
				com.anji.hyperneat.nd.NDFloatArray.MatrixIterator it = bias1.iterator();
				it.hasNext();
				it.next()
			) {
			it.set(0.20f);
		}
		NDFloatArray bias2 = new NDFloatArray(2, 2, 1);
		for (
				com.anji.hyperneat.nd.NDFloatArray.MatrixIterator it = bias2.iterator();
				it.hasNext();
				it.next()
			) {
			it.set(0.20f);
		}
		
		ActivatorND net = new GridNetND(allLayers, new NDFloatArray[] {wgts1, wgts2}, new NDFloatArray[] {bias1, bias2}, activationFn, 3, 0, true, "test" );
		
		NDFloatArray ins = new NDFloatArray(2, 2, 2);
		for (int z = 0; z < ins.getSizeOfDimension(0); z++) {
			for (int y = 0; y < ins.getSizeOfDimension(1); y++) {
				for (int x = 0; x < ins.getSizeOfDimension(2); x++) {
					if (z==y && y==x)
						ins.set(1.0f, x, y, z);
					else
						ins.set(0.0f, x, y, z);
				}
			}
		}
		
		NDFloatArray outs = net.next(ins);
		
		System.out.println(outs);
	}

	public GridNetND clone() {
		
	    ActivationFunction activationFunction = this.getActivationFunction();
	    NDFloatArray[] layers = new NDFloatArray[this.getLayers().length];
	    for (int i = 0; i < layers.length; i++) layers[i] = this.getLayers()[i].clone();
		NDFloatArray[] weights = new NDFloatArray[this.getLayers().length];
		for (int i = 0; i < weights.length; i++) weights[i] = this.weights[i].clone();
		NDFloatArray[] bias = null;
		if (enableBias) {
			bias = new NDFloatArray[this.bias.length];
			for (int i = 0; i < bias.length; i++) bias[i] = this.bias[i].clone();
		} 

		int maxDimensions = this.getMaxDimensions();
	    String name = new String(this.name);
	    int cyclesPerStep = this.getCyclesPerStep();
	    
	    GridNetND clone = new GridNetND(layers, weights, bias, activationFunction, maxDimensions, cyclesPerStep, enableBias, name);
	    
		return clone;
	}

	public void setMaxDimensions(int maxDimensions) {
		this.maxDimensions = maxDimensions;
	}

	public int getMaxDimensions() {
		return maxDimensions;
	}

	public void setActivationFunction(ActivationFunction activationFunction) {
		this.activationFunction = activationFunction;
	}

	public ActivationFunction getActivationFunction() {
		return activationFunction;
	}

	public void setLayers(NDFloatArray[] layers) {
		this.layers = layers;
	}

	public NDFloatArray[] getLayers() {
		return layers;
	}

	public void setLayer(int layerNum, NDFloatArray layer) {
		layers[layerNum] = layer;
	}
	
	public void setCyclesPerStep(int cyclesPerStep) {
		this.cyclesPerStep = cyclesPerStep;
	}

	public int getCyclesPerStep() {
		return cyclesPerStep;
	}
}
